#include "../DWMRICore/Scalar.h"
#include "../DWMRICore/VectorField.h"
#include "../DWMRICore/Fiber.h"

#include <string>

using namespace std;

char *pVectorFieldPathname = NULL;
char *pFiberPathname = NULL;
float stepSize = 0.1f;
float maxLength = 50.0f;
float faThreshold = 0.001f;
float angleThreshold = 0.707f;
int scale = 1;

CVectorField *pVectorField = NULL;

void PrintUsage()
{
	printf("Usage: ");
	printf("FiberTracking -vec <VectorField> -fiber <Fiber> [-stepsize <stepsize> -maxlength <maxlength> -fa_thresold <fa_threshold> -angle_threshold <fa_threshold> -scale <scale>]\n");
	printf("\n");
	printf("-vec <VectorField> = input VectorField nrrd file\n");
	printf("-fiber <Fiber> = output Fiber configure file\n");
	printf("-stepsize <stepsize> = step size\n");
	printf("-maxlength <maxlength> = maximum fiber length\n");
	printf("-fa_threshold <fa_thresold> = FA threshold used to differ white matter and grey matter\n");
	printf("-angle_threshold <angle_thresold> = angle threshold between two directions, in cosine\n");
	printf("-scale <scale> = scale of the result\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-vec") == 0) {
			pVectorFieldPathname = argv[++i];
		} else if (strcmp(argv[i], "-fiber") == 0) {
			pFiberPathname = argv[++i];
		} else if (strcmp(argv[i], "-stepsize") == 0) {
			stepSize = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-maxlength") == 0) {
			maxLength = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-fa_threshold") == 0) {
			faThreshold = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-angle_threshold") == 0) {
			angleThreshold = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-scale") == 0) {
			scale = (uint)atoi(argv[++i]);
		} else {
            printf("unknown %s\n", argv[i]);
			return 1;		// error
		}
	}

	if (pVectorFieldPathname == NULL || pFiberPathname == NULL)
		return -2;

	return 0;
}


int GetNextDirection(CPoint3F *pos, CVector3F *in, CVector3F *out)
{
	CVector3F vecs[3];
	int count = pVectorField->GetVectors(pos, vecs);
	if (count == 0) {
		return 0;
	}
	
	/* find best output direction */
	float max_v = 0;
	for (int i = 0; i < count; ++i) {
		float v = InnerProduct(in, &(vecs[i]));
		if (fabs(v) > fabs(max_v)) {
			max_v = v;
			out->m_x = vecs[i].m_x;
			out->m_y = vecs[i].m_y;
			out->m_z = vecs[i].m_z;
		}
	}

	if (fabs(max_v) < angleThreshold) {
		return 0;
	}

	if (max_v < 0.0f) {
		out->m_x = -out->m_x;
		out->m_y = -out->m_y;
		out->m_z = -out->m_z;
	}

	return 1;
}

void Tracking(Seed &seed, CFiber *fiber)
{
	uint maxSteps = (uint)(maxLength / (stepSize * 10.0f));
	uint fnum, bnum;
	CPoint3F *pF = new CPoint3F[maxSteps];
	CPoint3F *pB = new CPoint3F[maxSteps];
	CPoint3F lastPos;
	CVector3F lastDir;

	/* tracking forward */
	lastPos = seed.pos;
	lastDir = seed.dir;
	pF[0] = lastPos;
	for (fnum = 1; fnum < maxSteps; ++fnum) {
		CPoint3F nextPos;
		CVector3F nextDir;

		bool flag = true;
		for (int step = 0; step < 10; ++step) {
			nextPos = lastPos + lastDir * stepSize;

			if (GetNextDirection(&nextPos, &lastDir, &nextDir) == 0) {
				flag = false;
				break;
			}
			lastDir = nextDir;
			lastPos = nextPos;
		}
		if (flag == false)
			break;
		
		pF[fnum] = nextPos;
		lastDir = nextDir;
		lastPos = nextPos;
	}

	/* tracking backward */
	lastPos = seed.pos;
	lastDir = seed.dir;
	pB[0] = lastPos;
	for (bnum = 1; bnum < maxSteps; ++bnum) {
		CPoint3F nextPos;
		CVector3F nextDir;

		bool flag = true;
		for (int step = 0; step < 10; ++step) {
			nextPos = lastPos - lastDir * stepSize;

			if (GetNextDirection(&nextPos, &lastDir, &nextDir) == 0) {
				flag = false;
				break;
			}
			lastDir = nextDir;
			lastPos = nextPos;
		}
		if (flag == false)
			break;
		
		pB[bnum] = nextPos;
		lastDir = nextDir;
		lastPos = nextPos;
	}

	fiber->CreateFiber(pF, pB, fnum, bnum, &seed);


	SafeDeleteArray(pF);
	SafeDeleteArray(pB);
}

void FiberTracking()
{
	uint w = pVectorField->GetWidth();
	uint h = pVectorField->GetHeight();
	uint d = pVectorField->GetDepth();

	uint sw = w * scale;
	uint sh = h * scale;
	uint sd = d * scale;

	FILE *fp = fopen(pFiberPathname, "w");
	fprintf(fp, "%d %d %d %d %f %f\n", sw, sh, sd, scale, stepSize, maxLength);

	std::string pathname = pFiberPathname;
	size_t pos = pathname.find_last_of(SLASH);
	std::string path = pathname.substr(0, pos+1);

	for (uint z = 202; z < sd; ++z) {
        printf("Fiber tracking on slice z = %d ... ", z);
		CFiber **fibers = new CFiber*[sw*sh];
		int *count = new int[sw*sh];
		CPoint3F *seeds = new CPoint3F[sw*sh];
		
		/* Generate seeds */
		for (uint y = 0; y < sh; ++y) {
			for (uint x = 0; x < sw; ++x) {
				uint index = y * sw + x;
				seeds[index].m_x = (float)x/(float)scale + 1.0f / ((float)scale * 2.0f);
				seeds[index].m_y = (float)y/(float)scale + 1.0f / ((float)scale * 2.0f);
				seeds[index].m_z = (float)z/(float)scale + 1.0f / ((float)scale * 2.0f);
			}
		}

		/* tracking fibers for each seed */
		for (uint i = 0; i < sh*sw; ++i) {
			CVector3F vecs[3];
			count[i] = pVectorField->GetVectors(seeds[i], vecs);
			if (count[i] == 0) {
				fibers[i] = new CFiber[1];
				Seed temp;
				temp.dir = CVector3F(0.0f, 0.0f, 0.0f);
				temp.pos = seeds[i];
				fibers[i][0].CreateFiber(&(seeds[i]), &(seeds[i]), 1, 1, &temp);
				count[i] = 1;
			} else {
				fibers[i] = new CFiber[count[i]];
				for (int j = 0; j < count[i]; ++j) {
					Seed temp;
					temp.dir = vecs[j];
					temp.pos = seeds[i];
					Tracking(temp, &(fibers[i][j]));
				}
			}
		}

		/* Save fibers */
		//SaveFibersAsVTK("test.vtk", fibers, count, sw, sh);
      
		char confPathname[PATHNAME_SIZE];
		char vertPathname[PATHNAME_SIZE];
		sprintf(confPathname, "%sslice_z=%08.04f.conf", path.c_str(), (float)z/(float)scale);
		sprintf(vertPathname, "%sslice_z=%08.04f.vert", path.c_str(), (float)z/(float)scale);
		fprintf(fp, "%s\t%s\n", confPathname, vertPathname);

        SaveFibers(confPathname, vertPathname, fibers, count, sw, sh);

		for (uint i = 0; i < sh*sw; ++i) {
			SafeDeleteArray(fibers[i]);
		}
		SafeDeleteArray(fibers);
		SafeDeleteArray(count);

        printf("Done.\n");
	}

	fclose(fp);
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	pVectorField = new CVectorField();
	pVectorField->OpenVectorFieldFile(pVectorFieldPathname);

	FiberTracking();

	SafeDelete(pVectorField);
	return 0;
}
